# 只使用大模型进行转译，多线程，能够结合.h和.c文件同时转译，在log文件夹(自动生成)下记录报错日志

import os
import re
import subprocess
from zhipuai import ZhipuAI
from concurrent.futures import ThreadPoolExecutor, as_completed

# 配置智谱清言API密钥
with open('./Tool/api_key.txt', 'r') as api_key_file:
    API_KEY = api_key_file.read().strip()

SUCCESS = 1
FAIL = 0

MODEL_NAME = "glm-4-plus"  # 使用的模型名称
client = ZhipuAI(api_key=API_KEY)

MAX_THREAD = 10         # Thread number.
MAX_ITERATION_TIME = 3  # Iteration time of rewriting rust code with LLM. If exceed it, the iteration will stop

# 定义输入和输出文件夹
# INPUT_FOLDER = "Input/01-Primary/src" 
# OUTPUT_FOLDER = "Output/src"
INPUT_FOLDER = "Input/01-Primary/test" 
OUTPUT_FOLDER = "Output/test"
OUTPUT_DEBUG_DIR = OUTPUT_FOLDER + '/debug'
LOG_FOLDER = "log/test"  # 用于存储日志的文件夹

# 创建输出和日志文件夹（如果不存在）
os.makedirs(OUTPUT_FOLDER, exist_ok=True)
os.makedirs(LOG_FOLDER, exist_ok=True)

def extract_rust_code(response_content):
    """
    从文本中提取 Rust 代码块内容。
    """
    match = re.search(r"```rust\n(.*?)\n```", response_content, re.DOTALL)
    if match:
        return match.group(1).strip()  # 提取代码块中的内容
    return ""

def read_c_and_h_files(input_folder, base_name):
    """
    读取指定文件夹中与特定 `.c` 文件相关的 `.c` 和 `.h` 文件。
    返回拼接后的代码字符串。
    """
    c_file_path = os.path.join(input_folder, base_name + ".c")
    h_file_path = os.path.join(input_folder, base_name + ".h")
    code_parts = []

    # 读取 .c 文件
    if os.path.exists(c_file_path):
        with open(c_file_path, "r", encoding="utf-8") as c_file:
            code_parts.append(c_file.read())

    # 读取 .h 文件
    if os.path.exists(h_file_path):
        with open(h_file_path, "r", encoding="utf-8") as h_file:
            code_parts.append(h_file.read())

    return "\n\n".join(code_parts)

def compile_rust_code(rust_code, log_file_path):
    """
    尝试使用 rustc 编译 Rust 代码。
    返回编译是否成功以及错误信息（如果有）。
    """
    temp_file_path = f"{log_file_path}_temp.rs"
    try:
        with open(temp_file_path, "w", encoding="utf-8") as temp_file:
            temp_file.write(rust_code)

        result = subprocess.run(
            ["rustc", "--crate-type=lib", temp_file_path, '--out-dir', OUTPUT_DEBUG_DIR],
            capture_output=True,
            text=True
        )
        with open(log_file_path, "w", encoding="utf-8") as log_file:
            if result.returncode != 0:
                log_file.write(result.stderr)
            return result.returncode == 0
    finally:
        if os.path.exists(temp_file_path):
            os.remove(temp_file_path)

def optimize_rust_code(c_code, initial_rust_code, log_file_path, max_iterations=MAX_ITERATION_TIME):
    """
    迭代优化 Rust 代码。
    """
    current_rust_code = initial_rust_code
    for iteration in range(max_iterations):
        # print(f"---- Iteration {iteration + 1} in progress... ----")
        success = compile_rust_code(current_rust_code, log_file_path)

        if success:
            print(f"Compilation succeeded on iteration {iteration + 1}.")
            return current_rust_code, SUCCESS  # 成功编译的代码

        # 提示模型修复编译错误
        with open(log_file_path, "r", encoding="utf-8") as log_file:
            compile_errors = log_file.read()

        messages = [
            {
                "role": "user",
                "content": f"以下是C代码：\n```c\n{c_code}\n```以及初步转译后的Rust代码：\n```rust\n{current_rust_code}\n```\n"
                           f"但编译时出现了以下错误，请修复这些问题并重新提供完整的rust代码：\n{compile_errors}"
                           f"要求：给出完整的Rust代码，\不要使用占位符;直接输出代码块，不要在代码块前后有任何其他文字,不要省略代码,不要加上任何源代码没有的其他代码，如举例说明等。"
            }
        ]
        try:
            response = client.chat.completions.create(
                model=MODEL_NAME,
                messages=messages
            )
            current_rust_code = extract_rust_code(response.choices[0].message.content)
        except Exception as e:
            print(f"Error during API call: {e}")
            break

    print("Reached maximum iterations without successful compilation.")
    return current_rust_code, FAIL  # 返回最后尝试的代码

def convert_c_to_rust(c_code):
    """
    调用智谱AI接口，将C代码转换为Rust代码。
    """
    messages = [
        {
            "role": "user",
            "content": f"按照以下要求，将C代码及相关头文件内容转换为Rust代码,给出完整的Rust代码，\
            不要使用占位符;直接输出代码块，不要在代码块前后有任何其他文字,不要省略代码\
            不要加上任何源代码没有的其他代码，如举例说明等。：\n```c\n{c_code}\n```"
        }
    ]
    try:
        print('Translating...')
        
        response = client.chat.completions.create(
            model=MODEL_NAME,
            messages=messages
        )
        rust_code = extract_rust_code(response.choices[0].message.content)
        print('Initial translation over. Waiting for optimization...')
        return rust_code
    except Exception as e:
        print(f"Error during API call: {e}")
        return ""

def process_file(file_name):
    """
    处理单个文件的转译和优化。
    """
    base_name = os.path.splitext(file_name)[0]  # 获取文件名的基础部分（无扩展名）
    input_code = read_c_and_h_files(INPUT_FOLDER, base_name)  # 读取 .c 和 .h 文件

    print(f"==== Translating file: {file_name} ====")

    # 初步转换
    initial_rust_code = convert_c_to_rust(input_code)

    # 设置日志文件路径
    log_file_path = os.path.join(LOG_FOLDER, f"{base_name}")

    # 迭代优化
    optimized_rust_code, status = optimize_rust_code(input_code, initial_rust_code, log_file_path)

    # 确定输出路径
    output_path = os.path.join(OUTPUT_FOLDER, base_name + ".rs")

    # 将最终的Rust代码写入输出文件夹
    if optimized_rust_code and status == SUCCESS:
        with open(output_path, "w", encoding="utf-8") as file:
            file.write(optimized_rust_code)
        print(f"[Success]  {file_name} -> {output_path}")
    else:
        print(f"[Failed]   {file_name} See log: {log_file_path}")



# 使用多线程并行处理
with ThreadPoolExecutor(max_workers=MAX_THREAD) as executor:
    futures = [executor.submit(process_file, file_name) for file_name in os.listdir(INPUT_FOLDER) if file_name.endswith(".c")]
    for future in as_completed(futures):
        # 捕获线程执行中的异常
        try:
            future.result()
        except Exception as e:
            print(f"Error in thread execution: {e}")
